import tkinter as tk
import tkinter.messagebox
from tkinter import filedialog
import mne
import pylsl
import threading
import time
import pygame
import numpy as np
class EEG_Replay:
    def __init__(self):
        '''
            设置窗体
        '''
        self.mainWin = tk.Tk()  # 创建窗体
        self.mainWin.geometry('380x90+500+500')  # 设置窗体
        self.mainWin.title('离线数据LSL回放')
        self.hint = tk.StringVar()
        self.hint.set('请点击[选择脑电文件]选择一个文件')
        self.buttonText = tk.StringVar()
        self.buttonText.set('开始回放')
        '''界面对象'''
        self.loadBtn=tk.Button(self.mainWin, text='打开脑电文件', width=20, height=2,
                  command=self.Load_Data)
        self.loadBtn.grid(row=0, column=0, padx='10', pady='10', sticky=tk.E)  # 绑定open_file方法
        self.replayBtn=tk.Button(self.mainWin, textvariable=self.buttonText, width=20, height=2,
                  command=self.Replay)
        self.replayBtn.grid(row=0, column=1, padx='10', pady='10', sticky=tk.E)  # 绑定open_file方法
        self.replayBtn.config(state=tk.DISABLED)
        self.label=tk.Label(self.mainWin, textvariable=self.hint, anchor='w', relief=tk.SUNKEN, bg='#9c9', width=50, padx='10',
                 pady='0',
                 font=("微软雅黑", 9))
        self.label.grid(row=1, column=0, columnspan=2)

        #   环境变量
        self.replaying = False
        self.dataID=0
        tk.mainloop()
    def Load_Data(self):
        file = filedialog.askopenfile(title="选择脑电文件",
                                      filetypes=[('FIF文件', '*.fif'), ("BP设备数据", "*.vhdr"), ('EDF文件', '*.edf'),
                                                 ('BDF文件', '*.bdf')],
                                      initialdir=None)
        try:
            file=file.name
            suffix = file.split('.')[-1]
        except:
            tkinter.messagebox.showinfo('提示', '没有选择文件')
            #return
        # 根据文件后缀名确定mne的构建格式
        if suffix == 'vhdr':
            raw = mne.io.read_raw_brainvision(file, preload=True)
        elif suffix == 'fif':
            raw = mne.io.read_raw_fif(file, preload=True)
        elif suffix == 'edf':
            raw = mne.io.read_raw_edf(file, preload=True)
        elif suffix == 'bdf':
            raw = mne.io.read_raw_bdf(file, preload=True)
        else:
            tkinter.messagebox.showinfo('提示', '文件格式错误')
            return
        #   读取原信号，去除STIM通道数据
        self.rawMat=np.delete(raw.get_data(),-1,axis=0).transpose()
        self.replayBtn.config(state=tk.NORMAL)
        self.hint.set('采样率%s|电极：%s' %
                      (str(raw.info['sfreq'])
                       , str(raw.info['ch_names'])
                       ))
        '''构建回放类'''
        self.freq=raw.info['sfreq']
        self.info = pylsl.stream_info('type', 'eeg', len(raw.info['ch_names'])-1, raw.info['sfreq'], pylsl.cf_float32, 'dsffwerwer');
        self.outlet = pylsl.stream_outlet(self.info)
    def Replay(self):
        if self.replaying:
            self.replaying=False
            time.sleep(0.01)
            self.loadBtn.config(state=tk.NORMAL)
            self.buttonText.set('开始回放')
        else:
            self.replaying = True
            self.loadBtn.config(state=tk.DISABLED)
            time.sleep(0.01)
            t=threading.Thread(target=self.t_Replay,args=())
            t.setDaemon(True)
            t.start()
            self.buttonText.set('停止回放')
            self.hint.set('正在回放！')
    def t_Replay(self):
        clock=pygame.time.Clock()
        dataTotal=len(self.rawMat)
        self.hint.set('正在回放数据...')
        while self.replaying:
            clock.tick(self.freq)
            self.dataID+=1
            if self.dataID==dataTotal-1:self.dataID=0
            mysample = pylsl.vectorf(self.rawMat[self.dataID])
            self.outlet.push_sample(mysample)
            ''' 判断Trigger'''
            if self.rawMat[self.dataID][-1]!=0:
                self.hint.set('正在回放数据...【当前Trigger--%d】'%self.rawMat[self.dataID][-1])
if __name__=='__main__':
    e=EEG_Replay()